package com.leetcode;

import com.leetcode.common.TreeNode;
import com.leetcode.util.BinaryTreeUtil;

/**
 * 230. 二叉搜索树中第K小的元素
 * TODO: 在没有parent指针的情况下如何找任意一个节点的前驱和后继
 *
 * @author fy
 * @date 2022/4/26 20:17
 */
public class Solution230 {

    public int kthSmallest(TreeNode root, int k) {
        TreeNode node = findMin(root);
        int res = node.val;
        for (int i = 0; i < k - 1; i++) {
            res = successor(root, res);
        }
        return res;
    }

    /**
     * 查找一个节点val的后继节点
     *
     * @return
     */
    private int successor(TreeNode root, int val) {
        TreeNode node = findNode(root, val);
        if (node == null) {
            throw new RuntimeException("wrong val");
        }
        if (node.right != null) {
            node = node.right;
            while (node.left != null) {
                node = node.left;
            }
            return node.val;
        } else {
            return findParent(root, val).val;
        }
    }

    private TreeNode findNode(TreeNode root, int val) {
        if (root == null) {
            return null;
        }
        if (root.val > val) {
            return findNode(root.left, val);
        } else if (root.val < val) {
            return findNode(root.right, val);
        } else {
            return root;
        }
    }

    private TreeNode findParent(TreeNode root, int val) {
        if (root == null) {
            return null;
        }
        if (root.val > val) {
            if (root.left != null && root.left.val == val) {
                return root;
            }
            return findParent(root.left, val);
        } else if (root.val < val) {
            if (root.right != null && root.right.val == val) {
                return root;
            }
            return findParent(root.right, val);
        } else {
            return root;
        }
    }

    private TreeNode findMin(TreeNode root) {
        while (root.left != null) {
            root = root.left;
        }
        return root;
    }

    public static void main(String[] args) {
        TreeNode root = BinaryTreeUtil.createTree(new Integer[]{3, 1, 4, null, 2});
        BinaryTreeUtil.printTree(root);

        int k = 3;
        int res = new Solution230().kthSmallest(root, k);
        System.out.println(res);
    }

}
